package org.pyq.agent;

import javassist.*;

import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author pengYongQiang
 * @date 2020/11/6 15:57
 */
public class JDBCDemo {
    String beginSrc = "coderead.mybatis.log.JdbcCommonCollects inst=coderead.mybatis.log.JdbcCommonCollects.INSTANCE;";
    String endSrc = "result=inst.proxyConnection((java.sql.Connection)result);";
    String[] drivers = new String[]{"com.mysql.jdbc.NonRegisteringDriver", "com.mysql.cj.jdbc.NonRegisteringDriver"};
    private static Map<ClassLoader, ClassPool> classPoolMap = new ConcurrentHashMap<ClassLoader, ClassPool>();

    public byte[] transform(ClassLoader loader, String className) throws Exception {
        if (Arrays.stream(drivers).noneMatch(s -> s.equals(className))) {
            return null;
        }

        CtClass ctclass = toCtClass(loader, className);
        AgentByteBuild byteLoade = new AgentByteBuild(className, loader, ctclass);
        CtMethod connectMethod = ctclass.getMethod("connect", "(Ljava/lang/String;Ljava/util/Properties;)Ljava/sql/Connection;");
        AgentByteBuild.MethodSrcBuild build = new AgentByteBuild.MethodSrcBuild();
        build.setBeginSrc(beginSrc);
        build.setEndSrc(endSrc);
        byteLoade.updateMethod(connectMethod, build);
        return byteLoade.toBytecote();
    }

    protected static CtClass toCtClass(ClassLoader loader, String className) throws NotFoundException {
        if (!classPoolMap.containsKey(loader)) {
            ClassPool classPool = new ClassPool();
            classPool.insertClassPath(new LoaderClassPath(loader));
            classPoolMap.put(loader, classPool);
        }
        ClassPool cp = classPoolMap.get(loader);
        className = className.replaceAll("/", ".");
        return cp.get(className);
    }
}
